from osgeo import gdal,ogr

import numpy as np
import datetime,math
import random
import os


class TrainDT():
    def __init__(self):
        self.stopflag = False

    def writeTiff(self,im_data, im_geotrans, im_proj, path):
        if 'int8' in im_data.dtype.name:
            datatype = gdal.GDT_Byte
        elif 'int16' in im_data.dtype.name:
            datatype = gdal.GDT_UInt16
        elif 'int32' in im_data.dtype.name:
            datatype = gdal.GDT_UInt16
        else:
            datatype = gdal.GDT_Float32
        if len(im_data.shape) == 3:
            im_height, im_width,im_bands  = im_data.shape
        elif len(im_data.shape) == 2:
            im_height, im_width = im_data.shape
            im_bands=1

            #创建文件
        driver = gdal.GetDriverByName("GTiff")
        dataset = driver.Create(path, int(im_width), int(im_height), int(im_bands), datatype)
        if(dataset!= None):
            dataset.SetGeoTransform(im_geotrans) #写入仿射变换参数
            dataset.SetProjection(im_proj) #写入投影
        if im_bands==1:
            dataset.GetRasterBand(1).WriteArray(im_data)
        else:
            for i in range(im_bands):
                dataset.GetRasterBand(i+1).WriteArray(im_data[:,:,i])
        dataset.FlushCache()
        return dataset


    def calc_dis(self,mask,geoTrans,proj,cols, rows,noDataValue=0):
        memdrv = gdal.GetDriverByName('MEM')
        dst_ds = memdrv.Create('', cols, rows, 1, gdal.GDT_Int16)
        dst_ds.SetGeoTransform(geoTrans)
        dst_ds.SetProjection(proj)
        srcBand = dst_ds.GetRasterBand(1)
        srcBand.WriteArray(mask)

        memdrv2 = gdal.GetDriverByName('MEM')
        prox_ds = memdrv2.Create('', cols, rows, 1, gdal.GDT_Int16)
        prox_ds.SetGeoTransform(geoTrans)
        prox_ds.SetProjection(proj)
        proxBand = prox_ds.GetRasterBand(1)
        proxBand.SetNoDataValue(noDataValue)

        opt_string = 'NODATA=' + str(noDataValue)
        options = [opt_string]

        gdal.ComputeProximity(srcBand, proxBand, options)

        memdrv3 = gdal.GetDriverByName('MEM')
        proxIn_ds = memdrv3.Create('', cols, rows, 1, gdal.GDT_Int16)
        proxIn_ds.SetGeoTransform(geoTrans)
        proxIn_ds.SetProjection(proj)
        proxInBand = proxIn_ds.GetRasterBand(1)
        proxInBand.SetNoDataValue(noDataValue)
        opt_string2 = 'VALUES=' + str(noDataValue)
        options = [opt_string, opt_string2]
        gdal.ComputeProximity(srcBand, proxInBand, options)

        proxIn = proxIn_ds.ReadAsArray()
        proxOut = prox_ds.ReadAsArray()
        proxTotal = proxIn.astype(float) - proxOut.astype(float)
        # clip array
        # writeTiff(proxTotal, geoTrans, proj_img, npDistFileName)
        del dst_ds, prox_ds, proxIn_ds
        return proxTotal

    def pix2prj(self,gt, col, row):
        px = gt[0] + col * gt[1] + row  * gt[2]
        py = gt[3] + col * gt[4] + row * gt[5]
        return px, py

    def prj2pix(self,gt, x, y):
        det = gt[1] * gt[5] - gt[2] * gt[4]
        col = int((gt[5] * (x - gt[0]) - gt[2] * (y - gt[3])) / det)
        row = int((-gt[4] * (x - gt[0]) + gt[1] * (y - gt[3])) / det)
        return row, col


    def get_rand_patch(self,patch_img, idx):
        random_transformation = idx
        if random_transformation == 0:  # reverse first dimension
            patch_img = patch_img[::-1, :]
        elif random_transformation == 1:  # reverse second dimension
            patch_img = patch_img[:, ::-1]
        elif random_transformation == 2:  # transpose(interchange) first and second dimensions
            patch_img = patch_img.swapaxes(0, 1)
        elif random_transformation == 3:
            patch_img = np.rot90(patch_img, 1)
        elif random_transformation == 4:
            patch_img = np.rot90(patch_img, 2)
        elif random_transformation == 5:
            patch_img = np.rot90(patch_img, 3)
        else:
            pass
        return patch_img

    #  tif裁剪（tif像素数据，裁剪边长）
    def TifCropping(self,fileName, base_name,blocksize, path):
        SideLength = int((1 - math.sqrt(0.5)) * blocksize / 2)
        #  裁剪链表
        #  列上图像块数目
        dataset = gdal.Open(fileName)
        if dataset == None:
            print(fileName + "文件无法打开")
            return 0
        #  获取仿射矩阵信息
        geotrans = dataset.GetGeoTransform()
        #  获取投影信息
        im_proj = dataset.GetProjection()
        ColumnNum = int((dataset.RasterXSize - SideLength * 2) / (blocksize - SideLength * 2))
        #  行上图像块数目
        RowNum = int((dataset.RasterYSize - SideLength * 2) / (blocksize - SideLength * 2))
        for i in range(ColumnNum):
            for j in range(RowNum):
                cropped = dataset.ReadAsArray(i * (blocksize - SideLength * 2), j * (blocksize - SideLength * 2), blocksize,
                                              blocksize)
                px, py = self.pix2prj(geotrans, i * (blocksize - SideLength * 2), j * (blocksize - SideLength * 2))
                geotrans_cut = (px, geotrans[1], geotrans[2], py, geotrans[4], geotrans[5])
                self.writeTiff(cropped, geotrans_cut, im_proj, path + r'\{}_{}_{}.tif'.format(base_name,i, j))
        #  考虑到行列会有剩余的情况，向前裁剪一行和一列
        #  向前裁剪最后一列
        for i in range(ColumnNum):
            cropped = dataset.ReadAsArray(i * (blocksize - SideLength * 2), dataset.RasterYSize - blocksize, blocksize,
                                          blocksize)
            px, py = self.pix2prj(geotrans, i * (blocksize - SideLength * 2), (dataset.RasterYSize - blocksize))
            geotrans_cut = (px, geotrans[1], geotrans[2], py, geotrans[4], geotrans[5])
            self.writeTiff(cropped, geotrans_cut, im_proj, path + r'\{}_{}_{}.tif'.format(base_name,i, RowNum))

        #  向前裁剪最后一行
        for j in range(RowNum):
            cropped = dataset.ReadAsArray(dataset.RasterXSize - blocksize, j * (blocksize - SideLength * 2), blocksize,
                                          blocksize)
            px, py = self.pix2prj(geotrans, dataset.RasterXSize - blocksize, j * (blocksize - SideLength * 2))
            geotrans_cut = (px, geotrans[1], geotrans[2], py, geotrans[4], geotrans[5])
            self.writeTiff(cropped, geotrans_cut, im_proj, path + r'\{}_{}_{}.tif'.format(base_name,ColumnNum, j))

        #  向前裁剪右下角

        cropped = dataset.ReadAsArray(dataset.RasterXSize - blocksize, dataset.RasterYSize - blocksize, blocksize,
                                      blocksize)
        px, py = self.pix2prj(geotrans, dataset.RasterXSize - blocksize, dataset.RasterYSize - blocksize)
        geotrans_cut = (px, geotrans[1], geotrans[2], py, geotrans[4], geotrans[5])
        self.writeTiff(cropped, geotrans_cut, im_proj, path + r'\{}_{}_{}.tif'.format(base_name,ColumnNum, RowNum))



    #  tif裁剪（tif像素数据，裁剪边长）
    def TifCropping_viads(self,dataset_list,blocksize, path, base_name):
        img_path = path[0]
        mask_path = path[1]
        img_dataset = dataset_list[0]
        mask_dataset = dataset_list[1]
        SideLength = int((1 - math.sqrt(1)) * blocksize / 2)
        #  获取仿射矩阵信息
        geotrans = img_dataset.GetGeoTransform()
        #  获取投影信息
        im_proj = img_dataset.GetProjection()
        ColumnNum = int((img_dataset.RasterXSize - SideLength * 2) / (blocksize - SideLength * 2))
        #  行上图像块数目
        RowNum = int((img_dataset.RasterYSize - SideLength * 2) / (blocksize - SideLength * 2))
        idx = 0
        for i in range(ColumnNum):
            for j in range(RowNum):
                idx = idx + 1
                img_cropped = img_dataset.ReadAsArray(i * (blocksize - SideLength * 2), j * (blocksize - SideLength * 2), blocksize,
                                              blocksize)
                mask_cropped = mask_dataset.ReadAsArray(i * (blocksize - SideLength * 2), j * (blocksize - SideLength * 2),
                                                      blocksize,blocksize)
                # if mask_cropped.sum()==0:
                #     continue
                px, py = self.pix2prj(geotrans, i * (blocksize - SideLength * 2), j * (blocksize - SideLength * 2))
                geotrans_cut = (px, geotrans[1], geotrans[2], py, geotrans[4], geotrans[5])
                if len(img_cropped.shape) == 3:
                    img_cropped = img_cropped.transpose([1, 2, 0])
                # for idx in range(7):
                idx = random.randint(0, 6)
                img_cropped = self.get_rand_patch(img_cropped, idx)
                mask_cropped = self.get_rand_patch(mask_cropped, idx)
                self.writeTiff(img_cropped, geotrans_cut, im_proj, img_path + r'//' + base_name + '{}_{}_{}.tif'.format(i, j, idx))
                self.writeTiff(mask_cropped, geotrans_cut, im_proj, mask_path + r'//' + base_name + '{}_{}_{}.tif'.format(i, j, idx))
            # self.TrainDT_pyqtSignal.emit(0.5*(i*RowNum+j+1)/(ColumnNum*RowNum))
        #  考虑到行列会有剩余的情况，向前裁剪一行和一列
        #  向前裁剪最后一列
        idx = 0
        for i in range(ColumnNum):
            idx = idx + 1
            img_cropped = img_dataset.ReadAsArray(i * (blocksize - SideLength * 2), img_dataset.RasterYSize - blocksize, blocksize,
                                          blocksize)
            mask_cropped = mask_dataset.ReadAsArray(i * (blocksize - SideLength * 2), img_dataset.RasterYSize - blocksize,
                                                  blocksize,blocksize)
            # if mask_cropped.sum() == 0:
            #     continue
            px, py = self.pix2prj(geotrans, i * (blocksize - SideLength * 2), (img_dataset.RasterYSize - blocksize))
            geotrans_cut = (px, geotrans[1], geotrans[2], py, geotrans[4], geotrans[5])
            if len(img_cropped.shape) == 3:
                img_cropped = img_cropped.transpose([1, 2, 0])
            # for idx in range(7):
            idx = random.randint(0, 6)
            img_cropped = self.get_rand_patch(img_cropped, idx)
            mask_cropped = self.get_rand_patch(mask_cropped, idx)
            self.writeTiff(img_cropped, geotrans_cut, im_proj, img_path + r'//' + base_name + '{}_{}_{}.tif'.format(i, RowNum, idx))
            self.writeTiff(mask_cropped, geotrans_cut, im_proj, mask_path + r'//' + base_name + '{}_{}_{}.tif'.format(i, RowNum, idx))
        # self.TrainDT_pyqtSignal.emit(0.75)
        #  向前裁剪最后一行
        idx = 0
        for j in range(RowNum):
            idx = idx + 1
            img_cropped = img_dataset.ReadAsArray(img_dataset.RasterXSize - blocksize, j * (blocksize - SideLength * 2), blocksize,
                                          blocksize)
            mask_cropped = mask_dataset.ReadAsArray(img_dataset.RasterXSize - blocksize, j * (blocksize - SideLength * 2),
                                                  blocksize, blocksize)
            # if mask_cropped.sum() == 0:
            #     continue
            px, py = self.pix2prj(geotrans, img_dataset.RasterXSize - blocksize, j * (blocksize - SideLength * 2))
            geotrans_cut = (px, geotrans[1], geotrans[2], py, geotrans[4], geotrans[5])
            if len(img_cropped.shape) == 3:
                img_cropped = img_cropped.transpose([1, 2, 0])
            # for idx in range(7):
            idx = random.randint(0, 6)
            img_cropped = self.get_rand_patch(img_cropped, idx)
            mask_cropped = self.get_rand_patch(mask_cropped, idx)
            self.writeTiff(img_cropped, geotrans_cut, im_proj,
                      img_path + r'//' + base_name + '{}_{}_{}.tif'.format(ColumnNum, j, idx))
            self.writeTiff(mask_cropped, geotrans_cut, im_proj,
                      mask_path + r'//' + base_name + '{}_{}_{}.tif'.format(ColumnNum, j, idx))
        #  向前裁剪右下角
        img_cropped = img_dataset.ReadAsArray(img_dataset.RasterXSize - blocksize, img_dataset.RasterYSize - blocksize, blocksize,
                                      blocksize)

        mask_cropped = mask_dataset.ReadAsArray(img_dataset.RasterXSize - blocksize, img_dataset.RasterYSize - blocksize,
                                              blocksize, blocksize)
        # if mask_cropped.sum() == 0:
        #     return 0
        px, py = self.pix2prj(geotrans, img_dataset.RasterXSize - blocksize, img_dataset.RasterYSize - blocksize)
        geotrans_cut = (px, geotrans[1], geotrans[2], py, geotrans[4], geotrans[5])
        if len(img_cropped.shape) == 3:
            img_cropped = img_cropped.transpose([1, 2, 0])
        # for idx in range(7):
        idx = random.randint(0, 6)
        img_cropped = self.get_rand_patch(img_cropped, idx)
        mask_cropped = self.get_rand_patch(mask_cropped, idx)
        self.writeTiff(img_cropped, geotrans_cut, im_proj,
                  img_path + r'//' + base_name + '{}_{}_{}.tif'.format(ColumnNum, RowNum, idx))
        self.writeTiff(mask_cropped, geotrans_cut, im_proj,
                      mask_path + r'//' + base_name + '{}_{}_{}.tif'.format(ColumnNum, RowNum, idx))


    def save_data(self,in_file,save_img_path,save_mask_path,shp_path,blocksize=512,dis_val = 20,flag=0):
        # print('in_file====',in_file)
        # print('shp_path===',shp_path)
        in_ds = gdal.Open(in_file)
        geo_transform = in_ds.GetGeoTransform()
        proj_tuple = in_ds.GetProjection()
        cols = in_ds.RasterXSize  # 列数
        rows = in_ds.RasterYSize  # 行数
        bands = in_ds.RasterCount
        print(cols,rows,bands)
        nm = r'Z:\mask.tif'
        mem = gdal.GetDriverByName('GTiff')
        if os.path.exists(nm):
            os.remove(nm)
        out_mask_ds = mem.Create(nm, cols, rows,bands, gdal.GDT_Float32)
        out_mask_ds.SetGeoTransform(geo_transform)
        out_mask_ds.SetProjection(proj_tuple)
        m_ds = ogr.Open(shp_path)
        m_layer = m_ds.GetLayer()
        gdal.RasterizeLayer(out_mask_ds, [1], m_layer)  # 跟shp字段给栅格像元赋值
        if bands > 1:
            mask = out_mask_ds.ReadAsArray()
            mask = mask.sum(axis=0) > 0
        else:
            mask = out_mask_ds.ReadAsArray()
            mask = mask > 0


        clip_transform = out_mask_ds.GetGeoTransform()
        clip_proj = out_mask_ds.GetProjection()
        clip_row = out_mask_ds.RasterYSize
        clip_col = out_mask_ds.RasterXSize

        col_left, row_top = self.prj2pix(geo_transform, clip_transform[0], clip_transform[3])
        img_clip = in_ds.ReadAsArray(col_left, row_top, clip_col, clip_row)
        if os.path.exists(r'Z:\temp_img.tif'):
            os.remove(r'Z:\temp_img.tif')
        clip_img_ds = self.writeTiff(img_clip.transpose([1,2,0]), clip_transform, clip_proj, r'Z:\temp_img.tif')

        base_name = '{0:%Y%m%d%H%M%S%f}'.format(datetime.datetime.now())
        if flag==0:
            result_data = mask.astype(np.int)
        else:
            result_data = self.calc_dis(mask, clip_transform, clip_proj, cols, rows, noDataValue=0)
            result_data[result_data>dis_val] = dis_val
            result_data[result_data < -dis_val] = -dis_val
            result_data = result_data+dis_val
        if os.path.exists( r'Z:\temp_dis_pro.tif'):
            os.remove( r'Z:\temp_dis_pro.tif')
        clip_dis_ds = self.writeTiff(result_data, clip_transform, clip_proj, r'Z:\temp_dis_pro.tif')
        self.TifCropping_viads([clip_img_ds,clip_dis_ds], blocksize, [save_img_path,save_mask_path], base_name)
        # TifCropping_viads(clip_dis_ds, blocksize, save_mask_path, base_name)
        # TifCropping_viads(clip_img_ds, blocksize, save_img_path, base_name)
        del in_ds, out_mask_ds, clip_img_ds,clip_dis_ds
        os.remove(r'Z:\temp_img.tif')
        os.remove(r'Z:\temp_dis_pro.tif')
        os.remove(r'Z:\mask.tif')



if __name__ =='__main__':

    save_img_path = r'C:\Users\vipuser\Desktop\MySolution\data\rawdata\img'
    save_mask_path = r'C:\Users\vipuser\Desktop\MySolution\data\rawdata\mask'

    if not os.path.exists(save_img_path):
        os.makedirs(save_img_path)
    if not os.path.exists(save_mask_path):
        os.makedirs(save_mask_path)

    tdt = TrainDT()
    for i in range(23):
        in_file = r'C:\Users\vipuser\Desktop\MySolution\data\rawdata\TDT_data\building{}.tif'.format(i+1)
        shp_path = r'C:\Users\vipuser\Desktop\MySolution\data\rawdata\TDT_data\building_shp{}.shp'.format(i+1)
        tdt.save_data(in_file,save_img_path,save_mask_path,shp_path,blocksize=256)